import pandas as pd
import statsmodels.api as sm
import numpy as np

df = pd.read_csv('合并RID(填充缺失值).csv', encoding='gbk', engine='python')
Xcols = df.columns.tolist()[10:]
morethan03 = ['PIB', 'AV45', 'FBB', 'DIGITSCOR', 'DIGITSCOR_bl', 'MriType2', 'MOCA_bl', 'EcogPtMem_bl', 'EcogPtLang_bl', 'EcogPtVisspat_bl', 'EcogPtPlan_bl', 'EcogPtOrgan_bl', 'EcogPtDivatt_bl', 'EcogPtTotal_bl', 'EcogSPMem_bl', 'EcogSPLang_bl', 'EcogSPVisspat_bl', 'EcogSPPlan_bl', 'EcogSPOrgan_bl', 'EcogSPDivatt_bl', 'EcogSPTotal_bl', 'FDG_bl', 'PIB_bl', 'AV45_bl', 'FBB_bl']
nowNanCols = ['MOCA','EcogPtMem','EcogPtLang','EcogPtVisspat','EcogPtPlan','EcogPtOrgan','EcogPtDivatt','EcogPtTotal', 'update_stamp', 'TYPE2', 'EXAMDATE_bl', 'Years_bl', 'Month_bl', 'Month', 'M']
Xcols = list(set(Xcols) - set(morethan03) - set(nowNanCols))
typeCol = df['TYPE'].copy()

def modifyYcol(val):
    colName = 'TYPE'
    df[colName] = typeCol.copy()
    for index, row in df.iterrows():
        if row[colName] in val:
            df.loc[index, colName] = 1
        else:
            df.loc[index, colName] = 0

retTable=[]
def logit(XColName, YColName, YVal):
    Y = np.asarray(df[YColName]).astype(np.float64)
    X = np.asarray(df[XColName])
    mean = X.mean()
    std = X.std()
    X = (X-mean)/std
    logit = sm.Logit(Y, X)
    result = logit.fit()
    # 写进表里
    coef = result.params[0]
    p = result.pvalues[0]
    err = result.bse[0]
    row = {'feature':XColName, 'coef':coef, 'P(>|z|)':p, 'err':err, 'Yval':str(YVal)}
    retTable.append(row)

allType = [['CN'], ['AD'], ['LMCI', 'SMC', 'EMCI']]
nanList = []
for j in allType:
    modifyYcol(j)
    for i in Xcols:
        if np.issubdtype(df[i].dtype, np.number):
            if df[i].isnull().any():
                nanList.append(i)
            else:
                print('\nY:', j, 'X:', i)
                logit(i, 'TYPE', j)
print('NaN:', nanList)
pd.DataFrame(retTable).to_csv('逻辑回归结果.csv')